import numpy as np
import pandas as pd
import lightgbm as lgb
from scipy.special import expit
'''
    对外提供展示接口
'''
stock_file_name = 'input/stock_list.csv'
stock_list = pd.read_csv(stock_file_name)[['code','name']]
stock_list = stock_list.reset_index(drop=True)

def show_10_type1(predict_df):
    # predict_df = tmp_convert_boolean_int(predict_df)

    f_10_type1_df = filter_10cm_zt_1(predict_df)
    # X_predict = f_10_type1_df.drop(columns=['label', 'date','code','ret3'], errors='ignore')
    ret_columns = [col for col in f_10_type1_df.columns if col.startswith('ret')]
    X_predict = f_10_type1_df.drop(columns=['label', 'date', 'code'] + ret_columns, errors='ignore')
    
    prob_results = combined_probability_predict(
        model_recall_path='input/model/v3/v3_10_high_recall_zt_1.txt',
        model_precision_path='input/model/v3/v3_10_high_precision_zt_1.txt',
        X_data=X_predict
    )
    f_10_type1_final = f_10_type1_df.join(prob_results)
    f_10_type1_final = pd.merge(stock_list,f_10_type1_final,on='code')
    f_10_type1_final = f_10_type1_final[['date','code','name','r1','recall_prob','precision_prob','fused_prob','probability_label']]
    f_10_type1_final = f_10_type1_final[f_10_type1_final['probability_label'].isin(['极高概率', '高概率','中等概率'])]
    return f_10_type1_final


def show_w_ma(predict_df):
    f_10_type2_df = filter_w_ma(predict_df)
    X_predict = f_10_type2_df.drop(columns=['label', 'date','code','ret3'], errors='ignore')
    prob_results = combined_probability_predict(
        model_recall_path='input/model/v3/v3_h_recall_w_ma.txt',
        model_precision_path='input/model/v3/v3_h_precision_w_ma.txt',
        X_data=X_predict
    )
    # 合并原始数据与预测结果
    f_10_type2_final = f_10_type2_df.join(prob_results)
    f_10_type2_final = pd.merge(stock_list,f_10_type2_final,on='code')
    f_10_type2_final = f_10_type2_final[['date','code','name','r1','recall_prob','precision_prob','fused_prob','probability_label']]
    # f_10_type2_final = f_10_type2_final[f_10_type2_final['probability_label'].isin(['极高概率', '高概率','中等概率'])]
    return f_10_type2_final



def show_20_type2(predict_df):
    f_20_type2_df = f_20cm_4_9(predict_df)
    X_predict = f_20_type2_df.drop(columns=['label', 'date','code','ret3'], errors='ignore')
    prob_results = combined_probability_predict(
        model_recall_path='input/model/v3/v3_20_high_recall_4_9.txt',
        model_precision_path='input/model/v3/v3_20_high_precision_4_9.txt',
        X_data=X_predict
    )
    # 合并原始数据与预测结果
    f_20_type2_final = f_20_type2_df.join(prob_results)
    f_20_type2_final = pd.merge(stock_list,f_20_type2_final,on='code')
    f_20_type2_final = f_20_type2_final[['date','code','name','r1','recall_prob','precision_prob','fused_prob','probability_label']]
    f_20_type2_final = f_20_type2_final[f_20_type2_final['probability_label'].isin(['极高概率', '高概率','中等概率'])]
    return f_20_type2_final
    





'''
    10cm - type1 和 type2
'''
def filter_10cm_zt_1(train_df):
    f_l_df = filter_10cm(train_df)
    p_col(f_l_df)
    f_l_df = f_l_df.dropna(subset=['r15'])
    f_l_df = f_l_df[f_l_df['bias15'] >= 0]
    f_l_df = f_l_df[f_l_df['raise_buy'] < 4]
    f_l_df = f_l_df[f_l_df['is_zt'] & (f_l_df['zt_num'] <2)]
    f_l_df = f_l_df.reset_index(drop=True)
    p_col(f_l_df)
    return f_l_df

    
def filter_w_ma(train_df):
    f_l_df = train_df
    p_col(f_l_df)
    f_l_df = f_l_df.dropna(subset=['r15'])
    
    f_l_df = f_l_df[f_l_df['bias15'] >= 0]
    # f_l_df = f_l_df[f_l_df['zt_num'] < 1]
    f_l_df = f_l_df[(f_l_df['r1'] > -1 )&(f_l_df['r1'] < 7)]

    f_l_df = f_l_df[f_l_df['close'] >f_l_df['ma3w']]
    f_l_df = f_l_df[(f_l_df['h_w20'] >= - 3) &(f_l_df['h_w20'] <= 6)]

    f_l_df = f_l_df.reset_index(drop=True)
    p_col(f_l_df)
    return f_l_df
    


# def f_20cm_4_9(train_df):
#     f_l_df = filter_20cm(train_df)
#     p_col(f_l_df)
#     f_l_df = f_l_df.dropna(subset=['r15'])
#     f_l_df = f_l_df[f_l_df['bias15'] >= 0]
#     f_l_df = f_l_df[f_l_df['raise_buy'] < 4]
#     f_l_df = f_l_df[(f_l_df['r1'] >= 6) & (f_l_df['r1'] <=12)]
#     f_l_df = f_l_df[f_l_df['gt9_num'] <2]
#     f_l_df = f_l_df[(f_l_df['r10'] < 30)]
#     f_l_df = f_l_df.reset_index(drop=True)
#     p_col(f_l_df)
#     return f_l_df

    



def combined_probability_predict(model_recall_path, model_precision_path, X_data):
    """
    优化的双模型融合预测方案
    :param model_recall_path: 高召回模型路径 (召回率96%)
    :param model_precision_path: 高精准模型路径 (精准率58%)
    :param X_data: 待预测数据
    :return: 包含融合概率的结果DataFrame
    """
    # 加载模型
    model_recall = lgb.Booster(model_file=model_recall_path)
    model_precision = lgb.Booster(model_file=model_precision_path)
    
    # 阶段1：高召回模型预测
    recall_probs = model_recall.predict(X_data)
    recall_class3 = recall_probs[:, 3]  # 类别3概率
    
    # 阶段2：高精准模型预测
    precision_probs = model_precision.predict(X_data)
    precision_class3 = precision_probs[:, 3]  # 类别3概率
    
    # === 基于最新训练结果的融合策略 ===
    def dynamic_fusion(recall_prob, precision_prob):
        """动态融合函数 - 根据最新训练结果优化"""
        # 高召回模型置信度高时(>0.3)，优先考虑其预测
        if recall_prob > 0.3:
            # 高精准模型同时确认时，大幅提升概率
            if precision_prob > 0.4:
                return min(recall_prob * 1.8, 0.95)
            # 否则保持召回模型概率
            return recall_prob
        
        # 当高精准模型高度确信时(>0.5)，采用其预测
        if precision_prob > 0.5:
            return precision_prob * 1.2
        
        # 基础融合：加权平均（召回模型权重更高）
        recall_weight = 0.7  # 根据96%召回率调整权重
        precision_weight = 0.3  # 根据58%精准率调整权重
        
        # 加权融合
        fused_prob = (recall_weight * recall_prob + 
                      precision_weight * precision_prob)
        
        # 非线性放大（增强低概率区域的区分度）
        return np.power(fused_prob, 0.7)  # 指数小于1会放大低概率值
    
    # 应用融合策略
    fused_probs = np.array([
        dynamic_fusion(r, p) 
        for r, p in zip(recall_class3, precision_class3)
    ])
    
    # 概率归一化（确保在0-1范围内）
    fused_probs = np.clip(fused_probs, 0, 1)
    
    # 构建结果DataFrame
    results = pd.DataFrame({
        'recall_prob': recall_class3,
        'precision_prob': precision_class3,
        'fused_prob': fused_probs
    })
    
    # 添加概率解释标签
    conditions = [
        results['fused_prob'] > 0.6,
        results['fused_prob'] > 0.4,
        results['fused_prob'] > 0.2,
        results['fused_prob'] > 0.1
    ]
    labels = ['极高概率', '高概率', '中等概率', '低概率', '极低概率']
    results['probability_label'] = np.select(conditions, labels[:4], default=labels[4])
    
    return results



